############################################################################
############################################################################
##This file produces all figures for the simulation study in 
##conjoint example in Ratkovic and Tingley (2016).  It has 3 components:
## 1) Read in output.
## 2) Produce discovery plots.
## 3) Produces coverage plots.
##The file can be sourced and should be run without issue.
############################################################################
############################################################################


##Set to TRUE to print to directory with paper
forpaper<-FALSE

############################################################################
############################################################################
##Read in simulation output.
############################################################################
############################################################################

files.out<-list.files('Results_bch/')
files.out<-files.out[grep("n_",files.out)]		


##Declare some things.
sum1<-0
denom<-0
r.run<-NULL

n.run<-c(50,100,250,500,1000,2500)
p.run<-k.run<-c(76,181,356,531,706)
p.all<-NULL
zero.mat<-matrix(0,nr=length(p.run),nc=length(n.run))
rownames(zero.mat)<-p.run
colnames(zero.mat)<-n.run


DiscInterCover<-CoverInterNZ<-CoverInterZ<-RMSEInter<-FalseInter<-TrueInter<-RMSE<-RMSEZ<-RMSENZ<-denom.all<-denom.disc<-CoverDisc<-CoverNZ<-CoverZ<-FDR<-FPR<-FNR<-Rsq<-Time<-NULL


##Declare lists.
for(i.null in 1:16) DiscInterCover[[i.null]]<-CoverInterNZ[[i.null]]<-CoverInterZ[[i.null]]<-RMSEInter[[i.null]]<-RMSEZ[[i.null]]<-RMSE[[i.null]]<-RMSENZ[[i.null]]<-denom.all[[i.null]]<-denom.disc[[i.null]]<-CoverDisc[[i.null]]<-CoverZ[[i.null]]<-CoverNZ[[i.null]]<-FDR[[i.null]]<-FPR[[i.null]]<-FNR[[i.null]]<-Rsq[[i.null]]<-Time[[i.null]]<-FalseInter[[i.null]]<-TrueInter[[i.null]]<-zero.mat
	count.p<-1

##Loop through files.
for(i.file in files.out){
	load(paste("Results_bch/",i.file,sep=""))
	summary.out[is.na(summary.out)]<-0
	n.curr<-as.numeric(substr(i.file,3,3))
	k.curr<-as.numeric(substr(i.file,7,7))
	#print(dim(summary.out))
	p.all[count.p]<-summary.out["p",1]
	count.p<-count.p+1
	#print(summary.out["p",])
	for(i.method in 1:16){
		FDR[[i.method]][k.curr,n.curr] <- FDR[[i.method]][k.curr,n.curr] + summary.out["fdr",i.method]
		FPR[[i.method]][k.curr,n.curr] <- FPR[[i.method]][k.curr,n.curr]+summary.out["falsepos",i.method]
		FNR[[i.method]][k.curr,n.curr] <- FNR[[i.method]][k.curr,n.curr]+summary.out["falseneg",i.method]
		Rsq[[i.method]][k.curr,n.curr] <- Rsq[[i.method]][k.curr,n.curr]+summary.out["Rsq",i.method]
		CoverZ[[i.method]][k.curr,n.curr] <- CoverZ[[i.method]][k.curr,n.curr]+summary.out["cover0",i.method]
		CoverNZ[[i.method]][k.curr,n.curr] <- CoverNZ[[i.method]][k.curr,n.curr]+summary.out["covernz",i.method]
		CoverDisc[[i.method]][k.curr,n.curr] <- CoverDisc[[i.method]][k.curr,n.curr]+summary.out["coverdisc",i.method]
		RMSENZ[[i.method]][k.curr,n.curr] <- RMSENZ[[i.method]][k.curr,n.curr]+summary.out["rmse.nz",i.method]
		RMSE[[i.method]][k.curr,n.curr] <- RMSE[[i.method]][k.curr,n.curr]+summary.out["rmse.all",i.method]
		if(i.method==1) r.run<-c(r.run,(RMSE[[i.method]][k.curr,n.curr]+summary.out["rmse.all",i.method])
)
		RMSEZ[[i.method]][k.curr,n.curr] <- RMSEZ[[i.method]][k.curr,n.curr]+summary.out["rmse.zero",i.method]

		RMSEInter[[i.method]][k.curr,n.curr] <- RMSEZ[[i.method]][k.curr,n.curr]+summary.out["rmse.inters",i.method]

		FalseInter[[i.method]][k.curr,n.curr] <- FalseInter[[i.method]][k.curr,n.curr]+summary.out["falseinters",i.method]

		TrueInter[[i.method]][k.curr,n.curr] <- TrueInter[[i.method]][k.curr,n.curr]+summary.out["trueinters",i.method]
	#if(i.method==5) print(summary.out["trueinters",i.method])		

		DiscInterCover[[i.method]][k.curr,n.curr] <- DiscInterCover[[i.method]][k.curr,n.curr]+summary.out["discintercover",i.method]
		
		CoverInterNZ[[i.method]][k.curr,n.curr] <- CoverInterNZ[[i.method]][k.curr,n.curr]+summary.out["nzintercover",i.method]


		CoverInterZ[[i.method]][k.curr,n.curr] <- CoverInterZ[[i.method]][k.curr,n.curr]+summary.out["zintercover",i.method]

#CoverInterNZ<-CoverInterZ<-RMSEInter

		denom.disc[[i.method]][k.curr,n.curr] <- denom.disc[[i.method]][k.curr,n.curr]+summary.out["totdisc",i.method]
		denom.all[[i.method]][k.curr,n.curr] <- denom.all[[i.method]][k.curr,n.curr]+1	
	}

}


##Final cleanup.

	clean.denom<-function(x) {
		for(i in 1:16) {
			x[[i]]<-x[[i]]/denom.all[[i]]
			names(x)[i]<-colnames(summary.out)[i]
			#x[[i]][!is.finite(x[[i]])]<-0
			}
		x
		}


CoverInter<-CoverNZInter<-CoverZInter<-FDRInter<-FalseInter
for(i.fdr in 1:16){
	FDRInter[[i.fdr]]<-(FalseInter[[i.fdr]])/(FalseInter[[i.fdr]]+TrueInter[[i.fdr]])
	CoverZInter[[i.fdr]]<-CoverInterZ[[i.fdr]]/(FalseInter[[i.fdr]]+TrueInter[[i.fdr]])
	CoverNZInter[[i.fdr]]<-CoverInterNZ[[i.fdr]]/(FalseInter[[i.fdr]]+TrueInter[[i.fdr]])
	CoverInter[[i.fdr]]<-pmin((CoverInterZ[[i.fdr]]+CoverInterNZ[[i.fdr]])/(FalseInter[[i.fdr]]+TrueInter[[i.fdr]]),1)
}



	FDR<-clean.denom(FDR)
	FPR<-clean.denom(FPR)
	FNR<-clean.denom(FNR)
	Rsq<-clean.denom(Rsq)
	#Time<-clean.denom(Time)
	CoverZ<-clean.denom(CoverZ)
	CoverNZ<-clean.denom(CoverNZ)
	RMSE<-clean.denom(RMSE)
	RMSENZ<-clean.denom(RMSENZ)
	RMSEZ<-clean.denom(RMSEZ)
	RMSEInter<-clean.denom(RMSEInter)

	FalseInter<-clean.denom(FalseInter)
	TrueInter<-clean.denom(TrueInter)





		for(i.disc in 1:16) {
			CoverDisc[[i.disc]]<-CoverDisc[[i.disc]]/(denom.disc[[i.disc]])
			if(denom.disc[[i.disc]]<3) CoverDisc[[i.disc]]<-1/0
			#FDRInter[[i.disc]][!is.finite(FDRInter[[i.disc]])]<-0
			}
			names(CoverDisc)<-colnames(summary.out)
	

	TPR<-lapply(FNR,FUN=function(x) 12-x)


############################################################################
############################################################################
##Figures on effect discovery.
############################################################################
############################################################################



##Plotting function.

	lineplot.func.bw<-function(obj,plotset=c(1,11,14,7,16),zeroobj=TRUE,legend.loc="topright",plotleg=1,ylog.plot=FALSE,ab=FALSE,ylim.use=NULL){
		
	if(length(ylim.use)==0) for(i in plotset) ylim.use<-c(ylim.use,as.vector(obj[[i]]))
	ylim.use<-range(ylim.use,na.rm=TRUE)
	cols.all<-rep("black",16)
	lty.use<-rep(1,15)
	lty.use[plotset]<-c(1,2,3,5,6)[1:length(plotset)]
	for(i.k in 1:5){
		if(zeroobj) obj[[2]]<-obj[[2]]*0
		if(!ylog.plot) plot(n.run,obj[[i.k]][1,],log="x",type="n",ylim=ylim.use,xlab="",ylab="")
		if(ylog.plot) plot(n.run,obj[[i.k]][1,],log="xy",type="n",ylim=ylim.use,xlab="",ylab="")
		if(ab==TRUE) abline(h=.9,lwd=3,col=gray(.5))
		for(i.line in plotset) lines(n.run,obj[[i.line]][i.k,],lwd=ifelse(i.line%in%plotset[1:3],2,1),lty=lty.use[i.line])
	if(i.k%in%plotleg) legend(legend.loc,c("LASSOplus","LASSO","adaptive LASSO","Horseshoe","LASSO + OLS"),lwd=c(2.5,2.5,2.5,1,2),lty=c(1,2,3,5,6),bty="n",ncol=2,x.int=.1)
		}
	}


##PDFs.
	if(forpaper) {
		##Renaming figure for ease of reproduction
		#pdf("../../../paper/figs/Simplot_Disc.pdf",h=8*1.55,w=6*1.55)
		} else{
		##Renaming figure for ease of reproduction
		#pdf("Simplot_Disc.pdf",h=8*1.55,w=6*1.55)
		pdf("Figure2_SimulationDiscovery.pdf",h=8*1.55,w=6*1.55)

		}
	par(mfcol=c(5,3),mar=c(2,2,.1,.1),oma=c(2,2.5,2,0.2))
	lineplot.func.bw(FPR,ylim.use=c(0,15))
	lineplot.func.bw(FNR,ylim.use=c(0,15))
	lineplot.func.bw(FDR,ylim.use=c(0,1.05))
	
	mtext(outer=T,at=c(.19,.52,.86),c("Sample Size","Sample Size","Sample Size"),side=1,line=.3,cex=1)
	mtext(outer=T,at=c(.2,.5,.85),c("False Positives","False Negatives","False Discovery Rate"),side=3,line=-.0,cex=1.25)
	mtext(outer=T,at=c(.125+(0:4)*.2),c("706 Effects","531 Effects","356 Effects","181 Effects","76 Effects"),side=2,line=.5,cex=1)
	dev.off()
	

	if(forpaper){
		##Renaming figure for ease of reproduction
		#pdf("../../../paper/figs/Simplot_DiscInter.pdf",h=8*1.55,w=6*1.55)
	}else{
		##Renaming figure for ease of reproduction
		#pdf("Simplot_DiscInter.pdf",h=8*1.55,w=6*1.55)
		pdf("Figure3_SimulationInteractionDiscovery.pdf")
		}
	par(mfcol=c(5,3),mar=c(2,2,.1,.1),oma=c(2,2.5,2,0.2))
	lineplot.func.bw(FalseInter,ylim.use=c(0,12))
	TrueInter2<-TrueInter
	for(i in 1:16) TrueInter2[[i]]<-6-TrueInter[[i]]
	lineplot.func.bw(TrueInter2,ylim.use=c(0,12))
	lineplot.func.bw(FDRInter,ylim.use=c(0,1.05))
	
	mtext(outer=T,at=c(.19,.52,.86),c("Sample Size","Sample Size","Sample Size"),side=1,line=.3,cex=1)
	mtext(outer=T,at=c(.2,.5,.85),c("False Positives","False Negatives","False Discovery Rate"),side=3,line=-.0,cex=1.25)
	mtext(outer=T,at=c(.125+(0:4)*.2),c("706 Effects","531 Effects","356 Effects","181 Effects","76 Effects"),side=2,line=.5,cex=1)
	dev.off()

############################################################################
############################################################################
##Figures on effect discovery.
############################################################################
############################################################################

##Plotting function.
	coverplot.func.bw<-function(obj,plotset=c(1,11,14,7,16),legend.loc="topright",plotleg=1,lty.use=NULL){
	cols.all<-rep("black",16)
	lty.use<-rep(1,16)
	lty.use[plotset]<-c(1,2,3,5,6)[1:length(plotset)]
	for(i.k in 1:5){
		plot(n.run,obj[[i.k]][1,],log="x",type="n",ylim=c(0,1.3),xlab="",ylab="",xaxt="n")
		axis(1,at=n.run)
		abline(h=.9,lwd=2,col=gray(.5))
		for(i.line in 1:length(plotset)) lines(n.run,obj[[plotset[i.line]]][i.k,],lwd=ifelse(i.line<4,2,1),lty=lty.use[plotset[i.line]])
	if(i.k%in%plotleg) legend(legend.loc,c("LASSOplus","LASSO","adaptive LASSO","Horseshoe","LASSO + OLS"),lty=c(1,2,3,5,6),lwd=c(2.5,2.5,2.5,1),bty="n",ncol=2,x.int=.1,col=cols.all[plotset])
		
		}
	}
	

##Pdf
	if(forpaper){
		##Renaming figure for ease of reproduction
		#pdf("../../../paper/figs/Simplot_Cover.pdf",h=4.5*1.55*1.5*1.25,w=12)
	}else{
		##Renaming figure for ease of reproduction
		#pdf("Simplot_Cover.pdf",h=4.5*1.55*1.5*1.25,w=12)
		pdf("Figure4_SimulationCoverage.pdf",h=4.5*1.55*1.5*1.25,w=12)

		}
	par(mfcol=c(5,3),mar=c(2,2.5,.1,.1),oma=c(2,2.55,4,0.5))	
	coverplot.func.bw(CoverNZ)
	coverplot.func.bw(CoverDisc)	
	coverplot.func.bw(CoverInter)
	
	mtext(outer=T,at=c(.19,.52,.86),c("Sample Size","Sample Size","Sample Size"),side=1,line=.3,cex=1)
	mtext(outer=T,at=c(.2,.5,.85),c("Nonzero Effects","Discovered Effects","Discovered Interactions"),side=3,line=-.0,cex=1.25)
	mtext(outer=T,at=c(.125+(0:4)*.2),c("706 Effects","531 Effects","356 Effects","181 Effects","76 Effects"),side=2,line=.5,cex=1)
	dev.off()